import os
import sys
import glob
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker


# Fast computation of running mean, i.e., much faster than NumPy's convolve method:
def running_mean(data,window):
    cumsum = np.cumsum(np.insert(data,0,0))
    return (cumsum[window:] - cumsum[:-window]) / float(window)
    
    
def plot(path,algorithm,skipstart,skipstop):
   ext = glob.glob(path+"/*particlecount_"+algorithm+".csv")
   print(path,ext)
   filename = [os.path.basename(x) for x in sorted(ext)]

   
   # Absolute number of counted particles:
   n = []
   
   # Normalized number of counted particles:
   # (value - min) / (max - min)
   n_norm = []
   
   # Name of the activity performed:
   activity = []
   
   char1 = ""
   char2 = "_particlecount"
   
   with open(path+filename[0]) as f:
      print(filename[0])
      header = f.readlines()[:2]
      for i in range(0,len(header)):
         total_frames = int(header[0].split(":")[-1])
         fps = int(header[1].split(":")[-1])


   for f in range(0,len(filename)):
      void = np.loadtxt(path+filename[f],skiprows=skipstart)
      if skipstop > 0:
         n.append(running_mean(void[0:-skipstop],100))
      elif skipstop == 0:
         n.append(running_mean(void,100))
      else:
         print("Variable 'skipstop' must be >= 0. Exiting.")
         sys.exit(1)
      
      
      n_norm.append(np.divide(np.subtract(n[f],np.nanmin(n[f])),np.subtract(np.nanmax(n[f]),np.nanmin(n[f]))))
      
      activity.append(filename[f][filename[f].find(char1)+0:filename[f].find(char2)])
   
      print("Activity: {0:4s}   Mean: {1:.0f}   Stddev: {2:.0f}   Median: {3:.0f}   Range: {4:.0f}".format(activity[f],np.nanmean(n[f]),np.nanstd(n[f]),np.median(n[f]),(np.nanmax(n[f])-np.nanmin(n[f]))))
      
   # Create a time vector based on length of time series and the frame rate:
   t = []
   maxi = []
   t_maxi = []
   for f in range(0,len(n)):
      t.append(np.linspace(skipstart/fps,len(n[f])/fps,num=len(n[f]),endpoint=True))
      maxi.append(np.nanmax(n[f]))
      t_maxi.append(np.nanmax(t[f]))
   
   # Some plot settings:
   colours = ["red","black","darkblue","slategrey","orange","purple"]
   sym = ["o","+","v","^","x","D"]
   tick_spacing = 5
   xmin = 0
   xmax = (np.round(np.nanmax(t_maxi)/5)+1)*5
   ymin = 0
   ymax = np.round(np.nanmax(maxi)/100)*100 # np.nanmax(n)*1.1
   
   # Plot absolute number of counted particles:
   fig,ax = plt.subplots(1,1,)
   for f in range(0,len(filename)):
      ax.plot(t[f],n[f],color=colours[f],marker=None,linewidth=0.5,label=activity[f])
   ax.set_xlim(xmin,xmax)
   ax.set_ylim(ymin,ymax)
   ax.set_xlabel(r"Time / s")
   ax.set_ylabel(r"Number of counted particles / 1")
   ax.set_title(r"Video file: "+path+"*_subtracted_background_"+algorithm+".avi",fontsize=8)
   ax.text(xmax*0.75,ymax*0.71,"Frame rate: {0:3d} fps".format(fps),fontsize=8)
   plt.minorticks_on()
   ax.legend(loc="upper right",frameon=True,fontsize=8)
   plt.savefig(path+"All_activities_counted_particles_"+algorithm+".pdf",bbox_inches="tight")
   plt.close()
   
   
   #Some plot settings:
   ymax = 1.01
   
   # Plot normalized number of counted particles:
   fig,ax = plt.subplots(1,1,)
   for f in range(0,len(filename)):
      ax.plot(t[f],n_norm[f],color=colours[f],marker=None,linewidth=0.5,label=activity[f])
   ax.set_xlim(xmin,xmax)
   ax.set_ylim(ymin,ymax)
   ax.set_xlabel(r"Time / s")
   ax.set_ylabel(r"Normalized number of counted particles")
   ax.set_title(r"Video file: "+path+"*_subtracted_background_"+algorithm+".avi",fontsize=8)
   ax.text(xmax*0.75,ymax*0.71,"Frame rate: {0:3d} fps".format(fps),fontsize=8)
   plt.minorticks_on()
   ax.legend(loc="upper right",frameon=True,fontsize=8)
   plt.savefig(path+"All_activities_counted_particles_norm_"+algorithm+".pdf",bbox_inches="tight")
   plt.close()
   
   
if __name__ == "__main__":
   assert len(sys.argv) >= 1
   path = ""
   
   skipstart = 250
   skipstop = 250
   
   algorithm = "KNN"
   
   if len(sys.argv) == 4:
         plot(path+sys.argv[1]+"/"+sys.argv[2]+"/",sys.argv[3],algorithm,skipstart,skipstop)  
   elif len(sys.argv) == 3:
      plot(path+sys.argv[1]+"/",sys.argv[2],algorithm,skipstart,skipstop)
   elif path == "":
	   plot(sys.argv[1],algorithm,skipstart,skipstop)
   else:
      print("Wrong number of arguments. Exiting.") 
      sys.exit(1)

